# -*- coding: utf-8 -*-
"""watermark_mixin
批量给 PDF 添加水印线程功能模块。
请把 main.py 中 def add_watermarks_thread(...) 方法剪切粘贴到 WatermarkMixin 内。
"""

import os
import traceback
from tkinter import messagebox
import fitz  # PyMuPDF
from PIL import Image, ImageOps
import io


class WatermarkMixin:
    """PDF 添加水印相关方法 mixin"""

    def add_watermarks_thread(self, settings):
        """在后台线程中执行添加水印的操作 (使用 PIL 处理透明度)"""
        try:
            total_files = len(self.watermark_pdf_paths)
            processed_files = 0
            successful_files = 0
            failed_files_info = {}  # filename: error message

            # --- Pre-process watermark (convert to transparent PNG bytes) ---
            watermark_doc_src = None  # Source PDF or Image doc
            temp_doc_with_transparent_watermark = None  # Temp PDF holding the final transparent watermark
            watermark_rect = None

            try:
                # 1. Open source (PDF or Image)
                watermark_doc_src = fitz.open(self.watermark_file_path)
                if watermark_doc_src.page_count == 0:
                    raise ValueError(f"水印文件没有页面: {os.path.basename(self.watermark_file_path)}")
                watermark_page_src = watermark_doc_src[0]
                watermark_rect = watermark_page_src.rect

                # 2. Get Pixmap with Alpha channel
                watermark_pix = watermark_page_src.get_pixmap(alpha=True)

                # 3. Convert to PIL Image (RGBA)
                pil_img = Image.frombytes("RGBA", [watermark_pix.width, watermark_pix.height], watermark_pix.samples)

                # 4. Apply Opacity to Alpha Channel using PIL
                if settings["opacity"] < 1.0:  # Only apply if opacity is not 100%
                    try:
                        alpha = pil_img.split()[3]  # Get alpha channel
                        # Multiply existing alpha by the desired opacity factor
                        new_alpha = alpha.point(lambda p: int(p * settings["opacity"]))
                        pil_img.putalpha(new_alpha)  # Put the modified alpha back
                    except IndexError:
                        # Handle images without an alpha channel (treat as fully opaque base)
                        print(
                            f"水印文件 '{os.path.basename(self.watermark_file_path)}' 没有 Alpha 通道，将应用统一透明度。")
                        new_alpha_value = int(settings["opacity"] * 255)
                        alpha_layer = Image.new('L', pil_img.size, new_alpha_value)  # Create uniform alpha layer
                        pil_img.putalpha(alpha_layer)

                # 5. Convert Modified PIL Image back to PNG Bytes
                img_byte_arr = io.BytesIO()
                pil_img.save(img_byte_arr, format='PNG')
                final_transparent_bytes = img_byte_arr.getvalue()

                # 6. Create Temporary PDF containing the transparent watermark image
                temp_doc_with_transparent_watermark = fitz.open()
                temp_page = temp_doc_with_transparent_watermark.new_page(width=watermark_rect.width,
                                                                         height=watermark_rect.height)
                # Insert the image (which now has transparency baked in)
                temp_page.insert_image(watermark_rect, stream=final_transparent_bytes)

            except Exception as e:
                # Ensure source doc is closed if opened during failed preprocessing
                if watermark_doc_src:
                    watermark_doc_src.close()
                raise ValueError(f"无法加载或处理水印文件: {e}")  # Abort if watermark fails
            finally:
                # Close the original source doc now that we have the bytes/temp doc
                if watermark_doc_src:
                    watermark_doc_src.close()

            # --- Process each Target PDF ---
            if not temp_doc_with_transparent_watermark or watermark_rect is None:
                raise ValueError("水印预处理失败，无法继续。")

            for file_index, pdf_path in enumerate(self.watermark_pdf_paths):
                file_name = os.path.basename(pdf_path)
                base_name = os.path.splitext(file_name)[0]
                safe_name = "".join(c if c.isalnum() or c in (' ', '_', '-') else '_' for c in base_name).rstrip()
                if not safe_name: safe_name = f"file_{file_index}"

                output_pdf_path = os.path.join(settings["output_dir"], f"{safe_name}{settings['suffix']}.pdf")

                status_text = f"处理中: {file_name} ({file_index + 1}/{total_files})"
                progress = file_index / total_files
                self.after(0, self.update_progress_safe, progress, status_text)  # Update per file

                current_file_failed = False
                doc = None  # Target document
                try:
                    doc = fitz.open(pdf_path)

                    # Iterate through pages of the target document
                    for page_num in range(doc.page_count):
                        page = doc[page_num]
                        page_rect = page.rect

                        # --- Calculate watermark placement ---
                        wm_w = watermark_rect.width * settings["scale"]
                        wm_h = watermark_rect.height * settings["scale"]

                        # --- Apply Watermark using show_pdf_page with the pre-processed temp doc ---
                        if settings["position"] == "平铺":
                            step_x = wm_w * 1.1  # Example: 10% spacing
                            step_y = wm_h * 1.1
                            if step_x <= 0 or step_y <= 0: continue

                            for x in range(0, int(page_rect.width + step_x),
                                           int(step_x)):  # Extend range slightly for better coverage
                                for y in range(0, int(page_rect.height + step_y), int(step_y)):
                                    # Adjust position to center tiles better maybe? For now, simple grid.
                                    tile_rect = fitz.Rect(x, y, x + wm_w, y + wm_h)
                                    # Show the content from the temp doc (which has transparency)
                                    page.show_pdf_page(tile_rect, temp_doc_with_transparent_watermark, 0,
                                                       rotate=settings["rotation"],
                                                       overlay=settings["overlay"],
                                                       keep_proportion=True)
                        else:
                            # Single placement calculation (same as before)
                            pos_x, pos_y = 0, 0
                            # ... (calculate pos_x, pos_y based on settings["position"]) ...
                            if settings["position"] == "居中":
                                pos_x = (page_rect.width - wm_w) / 2
                                pos_y = (page_rect.height - wm_h) / 2
                            elif settings["position"] == "左上":
                                pos_x, pos_y = 0, 0
                            elif settings["position"] == "右上":
                                pos_x = page_rect.width - wm_w
                                pos_y = 0
                            elif settings["position"] == "左下":
                                pos_x = 0
                                pos_y = page_rect.height - wm_h
                            elif settings["position"] == "右下":
                                pos_x = page_rect.width - wm_w
                                pos_y = page_rect.height - wm_h

                            target_wm_rect = fitz.Rect(pos_x, pos_y, pos_x + wm_w, pos_y + wm_h)
                            if target_wm_rect.is_empty: continue

                            # Show the content from the temp doc (which has transparency)
                            page.show_pdf_page(target_wm_rect, temp_doc_with_transparent_watermark, 0,
                                               rotate=settings["rotation"],
                                               overlay=settings["overlay"],
                                               keep_proportion=True)

                    # Save the modified target document
                    doc.save(output_pdf_path, garbage=4, deflate=True, linear=True)
                    doc.close()
                    doc = None  # Indicate closed successfully

                except Exception as e:
                    print(f"处理文件 {file_name} 时出错: {e}")
                    print(traceback.format_exc())
                    current_file_failed = True
                    failed_files_info[file_name] = f"处理错误: {str(e)}"
                finally:
                    if doc is not None:  # Ensure target doc is closed even on error
                        doc.close()

                processed_files += 1
                if not current_file_failed:
                    successful_files += 1

            # --- Cleanup Temporary Watermark Doc ---
            if temp_doc_with_transparent_watermark is not None:
                temp_doc_with_transparent_watermark.close()

            # --- Final Report ---
            final_status = "水印添加完成"
            self.after(0, self.update_progress_safe, 1.0, final_status)

            # ... (Show final success/warning message - same as before) ...
            if failed_files_info:
                error_details = "\n失败文件详情:\n" + "\n".join(
                    [f"- {fname}: {emsg}" for fname, emsg in failed_files_info.items()])
                warning_message = (f"水印添加完成！\n\n成功: {successful_files} 个文件\n"
                                   f"失败: {len(failed_files_info)} 个文件\n"
                                   f"文件保存在: {settings['output_dir']}\n{error_details}")
                self.after(0, messagebox.showwarning, "部分完成", warning_message)
            else:
                success_message = (f"水印添加成功！\n\n已处理 {successful_files} 个文件\n"
                                   f"文件保存在: {settings['output_dir']}")
                self.after(0, messagebox.showinfo, "完成", success_message)


        except ValueError as ve:  # Catch setup errors (like loading watermark)
            self.after(0, messagebox.showerror, "错误", f"水印处理前检查出错: {str(ve)}")
            self.after(0, self.update_progress_safe, 0, "处理失败")
        except Exception as e:
            error_details = traceback.format_exc()
            self.after(0, self.show_error_message, f"添加水印过程中发生严重错误: {str(e)}\n{error_details}")
            self.after(0, self.update_progress_safe, 0, "处理失败")
        finally:
            # Ensure temp doc is closed even if outer try block fails after it's created
            if 'temp_doc_with_transparent_watermark' in locals() and temp_doc_with_transparent_watermark is not None:
                try:
                    temp_doc_with_transparent_watermark.close()
                except:
                    pass  # Ignore errors during cleanup close
            # Reset UI
            self.after(0, self.reset_ui_state, self.watermark_btn, self.watermark_progress_bar,
                       self.watermark_status_label, "准备就绪")
